import tkinter as tk
import cv2
import numpy as np 
from PIL import Image, ImageTk

from tkinter import messagebox
from tkinter.filedialog import askdirectory

import os

import sys
sys.path.append('./JetsonYolo')
from elements.yolo import OBJ_DETECTION

import cnn

simulate = True

import espeakng

object_classes = ['person',
                  'bicycle', 
                  'car', 
                  'motorcycle', 
                  'airplane', 
                  'bus', 
                  'train', 
                  'truck', 
                  'boat', 
                  'traffic light',
                  'fire hydrant', 
                  'stop sign', 
                  'parking meter', 
                  'bench', 
                  'bird', 
                  'cat', 
                  'dog', 
                  'horse', 
                  'sheep', 
                  'cow',
                  'elephant', 
                  'bear', 
                  'zebra', 
                  'giraffe', 
                  'backpack', 
                  'umbrella', 
                  'handbag', 
                  'tie', 
                  'suitcase', 
                  'frisbee',
                  'skis', 
                  'snowboard', 
                  'sports ball', 
                  'kite', 
                  'baseball bat', 
                  'baseball glove', 
                  'skateboard', 
                  'surfboard',
                  'tennis racket',
                  'bottle',
                  'wine glass',
                  'cup',
                  'fork',
                  'knife',
                  'spoon',
                  'bowl',
                  'banana',
                  'apple',
                  'sandwich',
                  'orange',
                  'broccoli',
                  'carrot', 
                  'hot dog', 
                  'pizza', 
                  'donut', 
                  'cake', 
                  'chair', 
                  'couch',
                  'potted plant', 
                  'bed', 
                  'dining table', 
                  'toilet', 
                  'tv', 
                  'laptop', 
                  'mouse', 
                  'remote', 
                  'keyboard', 
                  'cell phone',
                  'microwave',
                  'oven',
                  'toaster',
                  'sink',
                  'refrigerator',
                  'book',
                  'clock',
                  'vase',
                  'scissors',
                  'teddy bear',
                  'hair drier',
                  'toothbrush' ] 

from enum import Enum
class OperateMode(Enum):
    train = 1
    guide = 2

class ClassifySignalLightWindow(tk.Tk):
    labelPicWidth       = 640
    labelPicHeight      = 480
    buttonWidth         = 100
    buttonHeight        = 50
    textWidth           = 10
    textHeight          = 50
    tkWidth             = labelPicWidth
    tkHeigth            = labelPicHeight + buttonHeight * 4

    rect_color = (0, 0, 255)

    def __init__(self, *args, **kw):
        super().__init__()
        self.title("Guide")
        self.geometry(str(self.tkWidth) + "x" + str(self.tkHeigth))
        self.resizable(0, 0)
        
        self.object_detector = None
        self.cap = None
        self.operate_mode = OperateMode.train

        self.cnn_model = None

        self.speaker = espeakng.Speaker()

        self.my_model = None
        self.tf_saver = None
        self.tf_session = None

    def init_window(self):
        current_row = 0

        self.picture_label = tk.Label(self, text="Show Picture Area", font=("Arial", 24), bg="sky blue")
        self.picture_label.grid(row=current_row, columnspan=5)
        current_row = current_row + 1

        self.seperator_label_1 = tk.Label(self, text="", font=("Arial", 5), bg="sky blue")
        self.seperator_label_1.grid(row=current_row, columnspan=5)
        current_row = current_row + 1

        self.train_btn = tk.Radiobutton(self, text="Train", value=1, command=self.on_train_mode)
        self.train_btn.grid(row=current_row, column=1)
        self.guide_btn = tk.Radiobutton(self, text="Guide", value=2, command=self.on_guide_mode)
        self.guide_btn.grid(row=current_row, column=3)
        current_row = current_row + 1

        self.seperator_label_2 = tk.Label(self, text="", font=("Arial", 5), bg="sky blue")
        self.seperator_label_2.grid(row=current_row, columnspan=5)
        current_row = current_row + 1

        self.red_signal_btn = tk.Button(self, text="Red", bg="lightblue", width=10, command=self.on_red_signal)
        self.red_signal_btn.grid(row=current_row, column=0)
        self.green_signal_btn = tk.Button(self, text="Green", bg="lightblue", width=10, command=self.on_green_signal)
        self.green_signal_btn.grid(row=current_row, column=1)
        self.vehicle_signal_btn = tk.Button(self, text="Vehicle", bg="lightblue", width=10, command=self.on_vehicle_signal)
        self.vehicle_signal_btn.grid(row=current_row, column=2)

        self.prev_light_btn = tk.Button(self, text="Prev Light", bg="lightblue", width=10, command=self.on_prev_light)
        self.prev_light_btn.grid(row=current_row, column=3)
        self.next_light_btn = tk.Button(self, text="Next Light", bg="lightblue", width=10, command=self.on_next_light)
        self.next_light_btn.grid(row=current_row, column=4)
        current_row = current_row + 1

        self.seperator_label_3 = tk.Label(self, text="", font=("Arial", 5), bg="sky blue")
        self.seperator_label_3.grid(row=current_row, columnspan=5)

        self.select_folder_btn = tk.Button(self, text="Select Folder", bg="lightblue", width=10, command=self.on_select_folder)
        self.select_folder_btn.grid(row=current_row, column=0)

        self.prev_picture_btn = tk.Button(self, text="Prev Picture", bg="lightblue", width=10, command=self.on_prev_picture)
        self.prev_picture_btn.grid(row=current_row, column=2)

        self.next_picture_btn = tk.Button(self, text="Next Picture", bg="lightblue", width=10, command=self.on_next_picture)
        self.next_picture_btn.grid(row=current_row, column=4)
        current_row = current_row + 1

        if self.object_detector is None:
            self.object_detector = OBJ_DETECTION('weights/yolov5s.pt', object_classes) 

        if self.cnn_model is None:
            self.cnn_model = cnn.load_model()
 
        self.on_guide_mode()
        #self.on_train_mode()

    def on_train_mode(self):
        self.operate_mode = OperateMode.train

        self.red_signal_btn.config(state = 'normal')
        self.green_signal_btn.config(state = 'normal')
        self.vehicle_signal_btn.config(state = 'normal')

        self.prev_light_btn.config(state = 'normal')
        self.next_light_btn.config(state = 'normal')

        self.select_folder_btn.config(state = 'normal')
        self.prev_picture_btn.config(state = 'normal')
        self.next_picture_btn.config(state = 'normal')

        self.open_picture( '.', 'example.png' )

    def on_guide_mode(self):
        self.operate_mode = OperateMode.guide

        self.red_signal_btn.config(state = 'disabled')
        self.green_signal_btn.config(state = 'disabled')
        self.vehicle_signal_btn.config(state = 'disabled')

        self.prev_light_btn.config(state = 'disabled')
        self.next_light_btn.config(state = 'disabled')

        self.select_folder_btn.config(state = 'disabled')
        self.prev_picture_btn.config(state = 'disabled')
        self.next_picture_btn.config(state = 'disabled')
 
        tk.Message(self, text="Loading model...", width=100)

        if self.cap is None:
            if not simulate:
                self.cap = cv2.VideoCapture(0) 
            else:
                self.cap = cv2.VideoCapture('simulation.mp4')

        # calculate the scale and resized size
        ret, frame = self.cap.read() 
        s = frame.shape
        capWidth = s[1]
        capHeight = s[0]
        if capWidth <= self.labelPicWidth and capHeight <= self.labelPicHeight:
            self.scale = 1.0
            self.resizeWidth = capWidth
            self.resizeHeight = capHeight
        else: 
            widthScale = 1.0 * self.labelPicWidth / capWidth
            heightScale = 1.0 * self.labelPicHeight / capHeight
            self.scale = min(widthScale, heightScale)
            self.resizeWidth = int(capWidth * self.scale)
            self.resizeHeight = int(capHeight * self.scale)
     
        print("Capture Size:", capWidth, capHeight, self.scale)
        self.frame_count = 0
        self.run_guide()

    def save_signal_light(self, color):
        [(xmin,ymin), (xmax,ymax)] = self.current_object['bbox'] 
        xmin_disp = xmin - 5
        ymin_disp = ymin - 5
        xmax_disp = xmax + 5
        ymax_disp = ymax + 5
        cropped_image = self.image_cv[ymin:ymax, xmin:xmax]

        path = './image/'
        if color == cnn.TRAFFIC_LIGHT_RED:
            path = path + 'red'
        elif color == cnn.TRAFFIC_LIGHT_GREEN:
            path = path + 'green'
        elif color == cnn.TRAFFIC_LIGHT_VEHICLE:
            path = path + 'vehicle'
        else:
            pass


        (orig_filename, extension) = os.path.splitext( self.current_picture_file )
        signal_file_name = orig_filename + '_' + str(self.light_index) + '.png'

        if not os.path.exists(path):
            os.makedirs(path)

        if not os.path.exists(path):
            messagebox.showerror( "Error", "Cannot create path!" )
            
        path_name = path + '/' + signal_file_name
        #print( path_name )
        succ = cv2.imwrite( path_name, cropped_image )
        if not succ:
            messagebox.showerror( "Error", "Cannot write image!" )

    def announce_signal_color(self, color):
        mySpeaker = espeakng.Speaker()
        if color == cnn.TRAFFIC_LIGHT_RED:
            self.speaker.say('red light, please wait')
        elif color == cnn.TRAFFIC_LIGHT_GREEN:
            self.speaker.say('green light')
        elif color == cnn.TRAFFIC_LIGHT_VEHICLE:
            pass
        else:
            pass

    def on_red_signal(self):
        if self.operate_mode != OperateMode.train:
            return

        self.save_signal_light( cnn.TRAFFIC_LIGHT_RED )

        # proceed with the next signal
        self.find_next_light()
        self.disp_picture()

    def on_green_signal(self):
        if self.operate_mode != OperateMode.train:
            return

        self.save_signal_light( cnn.TRAFFIC_LIGHT_GREEN )

        # proceed with the next signal
        self.find_next_light()
        self.disp_picture()

    def on_vehicle_signal(self):
        if self.operate_mode != OperateMode.train:
            return

        self.save_signal_light( cnn.TRAFFIC_LIGHT_VEHICLE )

        # proceed with the next signal
        self.find_next_light()
        self.disp_picture()

        #self.speaker.say('yellow light, please wait')

    def on_prev_light(self):
        if self.operate_mode != OperateMode.train:
            return
 
        messagebox.showinfo( "Info", "Not implemented." )

    def find_next_light(self):
        while True:
            self.current_object = next( self.object_iter, None )
            if self.current_object is None:
                messagebox.showinfo( "Info", "This is the last traffic light!" )
                break

            label = self.current_object['label'] 
            score = self.current_object['score'] 

            if( label == "traffic light" ):
                # find a traffic light
                self.light_index = self.light_index + 1
                break
        
    def on_next_light(self):
        if self.operate_mode != OperateMode.train:
            return

        self.find_next_light()
        self.disp_picture()

    def on_select_folder(self):
        if self.operate_mode != OperateMode.train:
            return

        self.path = askdirectory()
        print( "path is %s" % self.path )

        self.picture_file_name = []
        for f in os.listdir(self.path):
            #print( "file name is: %s" % f )
            self.picture_file_name.append(f)

        self.picture_file_iter = iter( self.picture_file_name )

        self.on_next_picture()
 

    def on_prev_picture(self):
        if self.operate_mode != OperateMode.train:
            return

        messagebox.showinfo( "Info", "Not implemented." )

    def on_next_picture(self):
        if self.operate_mode != OperateMode.train:
            return

        while True:
            self.current_picture_file = next( self.picture_file_iter, None )
            if self.current_picture_file is None:
                messagebox.showinfo( "Info", "This is the last file!" )
                break
   
            rtn = self.open_picture( self.path, self.current_picture_file )
            if rtn:
                break

    def run_guide(self):
        #print( "refreshing...\n" )

        ret, self.frame = self.cap.read() 
        self.frame_count = self.frame_count + 1
        image_cv_with_obj = self.frame.copy()

        s = self.frame.shape
        capWidth = s[1]
        capHeight = s[0]

        signals = []
        objects = self.object_detector.detect( self.frame )
        for obj in objects:
            label = obj['label']
            score = obj['score']
            if( label == "traffic light" ):
                [(xmin,ymin), (xmax,ymax)] = obj['bbox'] 
                """
                if xmin > xmax:
                    xmin, xmax = xmax, xmin
                if ymin > ymax:
                    ymin, ymax = ymax, ymin
                xmin_disp = xmin - 5
                if xmin_disp < 0:
                    xmin_disp = 0
                ymin_disp = ymin - 5
                if ymin_disp < 0:
                    ymin_disp = 0
                xmax_disp = xmax + 5
                if xmax_disp > self.resizeWidth:
                    xmax_disp = self.resizeWidth
                ymax_disp = ymax + 5
                if ymax_disp > self.resizeHeight:
                    ymax_disp = self.resizeHeight
                """
                signals.append( obj )

                xmin = xmin - 5
                ymin = ymin - 5
                xmax = xmax + 5
                ymax = ymax + 5

                if xmin < 0 :
                    xmin = 0
                if xmax < 0:
                    xmax = 0
                if ymin < 0:
                    ymin = 0
                if ymax < 0:
                    ymin = 0
                if xmin > capWidth:
                    xmin = capWidth
                if xmax > capWidth:
                    xmax = capWidth
                if ymin > capHeight:
                    ymin = capHeight
                if ymax > capHeight:
                    ymax = capHeight

                signal_image = self.frame[ymin:ymax, xmin:xmax]

                print( "xmin = %d, ymin = %d, xmax = %d, ymax = %d" % (xmin, ymin, xmax, ymax) )
                disp_image = self.frame.copy()

                signal_light = cnn.predict( self.cnn_model, signal_image )
                print( "result is %s" % cnn.get_light_name(signal_light) )

                self.announce_signal_color( signal_light )

                """
                # display the image for debug
                imgCVRGB = cv2.cvtColor(disp_image, cv2.COLOR_BGR2RGB)
                img = Image.fromarray(imgCVRGB)
                self.image_tk = ImageTk.PhotoImage(image=img)
                self.picture_label.configure(image = self.image_tk, bg="pink")
                """

                # write to file to debug
                """
                path = './image/capture/'
                if not os.path.exists(path):
                    os.makedirs(path)
                path_name = "%s_frame_%08d_%02d.png" % (path, self.frame_count, len(signals) )  
                #print( path_name )
                succ = cv2.imwrite( path_name, signal_image )
                """
        

        # display the image for debug
        for i in range(len(signals)):
            current_object = signals[i]
            label = current_object['label'] 
            score = current_object['score'] 
            [(xmin,ymin), (xmax,ymax)] = current_object['bbox'] 
            #print( "label = %s, %d, %d, %d, %d" % (label, xmin, ymin, xmax, ymax) )
            xmin_disp = xmin - 5
            ymin_disp = ymin - 5
            xmax_disp = xmax + 5
            ymax_disp = ymax + 5
            image_cv_with_obj = cv2.rectangle(image_cv_with_obj, (xmin_disp,ymin_disp), (xmax_disp,ymax_disp), self.rect_color, 2) 

        # display picture with rectangle
        image_cv_with_obj = cv2.resize( image_cv_with_obj, (self.resizeWidth, self.resizeHeight), cv2.INTER_AREA )
        image_cv_rgb = cv2.cvtColor( image_cv_with_obj, cv2.COLOR_BGR2RGB)
        image_pil = Image.fromarray( image_cv_rgb )
        self.image_tk = ImageTk.PhotoImage(image=image_pil)
        self.picture_label.configure(image = self.image_tk, bg="pink")

        self.after(100, self.run_guide)

    def open_picture(self, path, file_name):

        path_name = path + '/' + file_name
        print( "file path is: %s" % path_name )

        self.image_cv = cv2.imread( path_name )

        if self.image_cv is None:
            print("Read Fail!")
            return False

        self.current_picture_file = file_name
        self.light_index = 1

        s = self.image_cv.shape
        picWidth = s[1]
        picHeight = s[0]
        # print("Picture Size:", picWidth, picHeight)
        if (picWidth > self.labelPicWidth) or (picHeight > self.labelPicHeight):
            widthScale = 1.0 * self.labelPicWidth / picWidth
            heightScale = 1.0 * self.labelPicHeight / picHeight
        
            scale = min(widthScale, heightScale)
        
            resizeWidth = int(picWidth*scale)
            resizeHeight = int(picHeight*scale)
        
            print( "scale = %d, width = %d, height=%d" %(scale, resizeWidth, resizeHeight) )
    
            self.image_cv = cv2.resize( self.image_cv, (resizeWidth, resizeHeight), cv2.INTER_AREA )

        self.detected_objects = self.object_detector.detect( self.image_cv )

        self.object_iter = iter( self.detected_objects )
        self.find_next_light()
        
        self.disp_picture()

        return True

    def disp_picture(self):
        image_cv_with_obj = self.image_cv.copy()

        # draw the rectangle
        if not self.current_object is None:
            label = self.current_object['label'] 
            score = self.current_object['score'] 
            [(xmin,ymin), (xmax,ymax)] = self.current_object['bbox'] 
            print( "label = %s, %d, %d, %d, %d" % (label, xmin, ymin, xmax, ymax) )
            xmin_disp = xmin - 5
            ymin_disp = ymin - 5
            xmax_disp = xmax + 5
            ymax_disp = ymax + 5
            image_cv_with_obj = cv2.rectangle(image_cv_with_obj, (xmin_disp,ymin_disp), (xmax_disp,ymax_disp), self.rect_color, 2) 

        # display picture with rectangle
        image_cv_rgb = cv2.cvtColor( image_cv_with_obj, cv2.COLOR_BGR2RGB)
        image_pil = Image.fromarray( image_cv_rgb )
        self.image_tk = ImageTk.PhotoImage(image=image_pil)
        self.picture_label.configure(image = self.image_tk, bg="pink")

c = ClassifySignalLightWindow()
c.init_window()
c.mainloop()
